

from gcl_configs import args
import os
import function_tool as ft
import gcl_model as gclm
import torch


def main(args):
    datasetname = args.data_name
    label_list = args.label_list
    graphname = args.graph_name
    device = args.device

    # get the path to save the results for gcl

    # data location
    # original_data_path = '../data100/{}/'.format(datasetname)
    original_data_path = '../marldata20231011/{}/'.format(datasetname)
    original_data_file = '{}.pkl'.format(graphname)
    transformed_data_filename = '{}_transform20231011.pkl'.format(graphname)
    pyg_graph_data_location = original_data_path + original_data_file

    # gcl model weights location
    # gcl model weights location
    gcl_model_weights_path = 'model_weights_marl20231011/gcl/{}/{}/'.format(datasetname, graphname)
    gcl_model_weights_file = '{}_model.pkl'.format(graphname)
    gcl_model_weights_location = gcl_model_weights_path + gcl_model_weights_file

    # get the gcl encoder
    input_dim, hidden_dim, num_layers = args.input_dim_encoder, args.hidden_dim_encoder, args.num_layers_encoder
    gcl_encoder_model, contrast_model, gcl_optimizer = gclm.model().get_model(input_dim, hidden_dim, num_layers, device)
    # needed
    gcl_encoder_model.load_state_dict(torch.load(gcl_model_weights_location))
    # select: 0 or 1

    # to non-star data with distance and new feature
    dataset = ft.load_nx_to_pyg_for_ls(pyg_graph_data_location, args.input_dim_encoder)
    _ = ft.load_data_with_transformed_feature(dataset, original_data_path, transformed_data_filename, gcl_encoder_model,
                                          device)
    print('Finished and transform into non-star data!')

if __name__ == '__main__':
    """
    datasetname = 'crg_gnp_random_graph'  # 'crg_gnp_random_graph', 'rpt_rt_tree_graph','rc_bg_graph'
    graphname = 'crg_gnp_0.2'  # 'crg_gnp_p' p=0.2~0.9. 'rpt_rt','rc_bg'
    label_list = ['crg', 'gnp']  # ['crg', 'gnp'], ['rpt', 'rt'], ['rc', 'bg']
    give the metric dimension of each graph
    """
    # ////////////////////////----------EXPERIMENT SETUP----------////////////////////////////
    # experiment
    args.data_name = 'rpt_rt_tree_graph'
    args.label_list = ['rpt', 'rt']
    args.graph_name = 'rpt_rt'
    main(args)

    # args.data_name = 'rc_bg_graph'
    # args.label_list = ['rc', 'bg']
    # args.graph_name = 'rc_bg'
    # ps = [0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9]
    # args.data_name = 'crg_gnp_random_graph'
    # args.label_list = ['crg', 'gnp']
    # for p in ps:
    #     #p=0.2
    #     args.graph_name = 'crg_gnp_{}'.format(p)
    #
    #     args.feature_dim = 64
    #     args.input_dim_encoder = args.feature_dim
    #     args.epochs_gcl = 100
    #
    #     main()
